package com.dms.modules.message.websocket;

import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.dms.modules.message.dto.BusinessMessageDTO;
import com.dms.modules.message.dto.MessageDTO;
import com.dms.modules.message.dto.MessageDetailDTO;
import com.dms.modules.message.service.BusinessMessageService;
import com.dms.modules.message.service.MessageService;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.google.gson.JsonSyntaxException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.context.annotation.Lazy;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.time.LocalDateTime;
import java.time.temporal.ChronoUnit;
import java.util.Collections;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.*;

@Slf4j
@Component
public class MessageWebSocket extends TextWebSocketHandler {

    private static final Map<Long, WebSocketSession> sessions = new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<Long, LocalDateTime> LAST_HEARTBEAT = new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<String, CompletableFuture<Boolean>> MESSAGE_CONFIRMATIONS = new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<Long, ConcurrentLinkedQueue<MessageDetailDTO>> OFFLINE_MESSAGES = new ConcurrentHashMap<>();
    
    private static BusinessMessageService businessMessageService;
    private static MessageService messageService;
    private static final Gson gson = new Gson();
    
    // 心跳超时时间（毫秒）
    private static final long HEARTBEAT_TIMEOUT = 30000;
    // 消息确认超时时间（毫秒）
    private static final long MESSAGE_CONFIRMATION_TIMEOUT = 5000;
    
    // 心跳检查调度器
    private static final ScheduledExecutorService heartbeatScheduler = Executors.newSingleThreadScheduledExecutor(r -> {
        Thread thread = new Thread(r, "websocket-heartbeat-checker");
        thread.setDaemon(true);
        return thread;
    });
    
    static {
        // 启动心跳检查任务
        heartbeatScheduler.scheduleAtFixedRate(() -> {
            try {
                checkHeartbeats();
            } catch (Exception e) {
                log.error("心跳检查任务异常", e);
            }
        }, 10, 10, TimeUnit.SECONDS);
    }

    @Autowired
    @Lazy
    public void setBusinessMessageService(BusinessMessageService service) {
        MessageWebSocket.businessMessageService = service;
    }

    @Autowired
    public void setMessageService(MessageService service) {
        MessageWebSocket.messageService = service;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        Long userId = getUserId(session);
        if (userId != null) {
            sessions.put(userId, session);
            LAST_HEARTBEAT.put(userId, LocalDateTime.now());
            log.info("WebSocket连接建立: userId={}", userId);
            
            // 发送未读消息
            sendUnreadMessages(userId);
            
            // 发送离线消息
            sendOfflineMessages(userId);
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        Long userId = getUserId(session);
        if (userId != null) {
            sessions.remove(userId);
            LAST_HEARTBEAT.remove(userId);
            log.info("WebSocket连接关闭: userId={}", userId);
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        if (!StringUtils.hasText(message.getPayload())) {
            return;
        }

        try {
            JsonObject jsonMessage = gson.fromJson(message.getPayload(), JsonObject.class);
            String type = jsonMessage.has("type") ? jsonMessage.get("type").getAsString() : null;
            
            if (!StringUtils.hasText(type)) {
                log.warn("消息类型为空: userId={}, message={}", getUserId(session), message.getPayload());
                return;
            }
            
            switch (type) {
                case "heartbeat":
                    handleHeartbeat(getUserId(session));
                    break;
                case "message":
                    handleBusinessMessage(getUserId(session), jsonMessage);
                    break;
                case "ack":
                    handleMessageAck(jsonMessage);
                    break;
                default:
                    log.warn("未知的消息类型: type={}, userId={}", type, getUserId(session));
            }
        } catch (JsonSyntaxException e) {
            log.error("消息格式错误: userId={}, message={}", getUserId(session), message.getPayload(), e);
        } catch (Exception e) {
            log.error("处理消息失败: userId={}, message={}", getUserId(session), message.getPayload(), e);
        }
    }

    private void handleHeartbeat(Long userId) {
        LAST_HEARTBEAT.put(userId, LocalDateTime.now());
        sendHeartbeatResponse(userId);
    }

    private void handleBusinessMessage(Long userId, JsonObject jsonMessage) {
        try {
            if (!jsonMessage.has("data") || !jsonMessage.has("messageId")) {
                log.warn("业务消息格式错误: userId={}, message={}", userId, jsonMessage);
                return;
            }

            String messageId = jsonMessage.get("messageId").getAsString();
            BusinessMessageDTO messageDTO = gson.fromJson(jsonMessage.get("data"), BusinessMessageDTO.class);
            
            if (messageDTO == null || messageDTO.getType() == null) {
                log.warn("业务消息数据为空: userId={}, messageId={}", userId, messageId);
                return;
            }

            // 验证消息所属用户
            if (!Objects.equals(messageDTO.getUserId(), userId)) {
                log.warn("消息用户ID不匹配: messageUserId={}, sessionUserId={}", messageDTO.getUserId(), userId);
                return;
            }

            businessMessageService.sendBusinessMessage(messageDTO);
            sendMessageAck(userId, messageId);
            
            log.info("业务消息处理成功: userId={}, messageId={}, type={}", userId, messageId, messageDTO.getType());
        } catch (Exception e) {
            log.error("处理业务消息失败: userId={}", userId, e);
        }
    }

    private void handleMessageAck(JsonObject jsonMessage) {
        if (!jsonMessage.has("messageId")) {
            return;
        }

        String messageId = jsonMessage.get("messageId").getAsString();
        CompletableFuture<Boolean> future = MESSAGE_CONFIRMATIONS.remove(messageId);
        if (future != null && !future.isDone()) {
            future.complete(true);
        }
    }

    public CompletableFuture<Boolean> sendMessage(Long userId, MessageDetailDTO message) {
        if (userId == null || message == null) {
            return CompletableFuture.completedFuture(false);
        }

        WebSocketSession session = sessions.get(userId);
        if (session != null && session.isOpen()) {
            try {
                String messageId = generateMessageId();
                JsonObject wrapper = new JsonObject();
                wrapper.addProperty("type", "message");
                wrapper.addProperty("messageId", messageId);
                wrapper.add("data", gson.toJsonTree(message));
                
                CompletableFuture<Boolean> confirmation = new CompletableFuture<>();
                MESSAGE_CONFIRMATIONS.put(messageId, confirmation);
                
                session.sendMessage(new TextMessage(gson.toJson(wrapper)));
                log.info("消息发送成功: userId={}, messageId={}", userId, messageId);
                
                // 设置确认超时
                scheduleConfirmationTimeout(messageId, confirmation);
                
                return confirmation;
            } catch (IOException e) {
                log.error("发送消息失败: userId={}, error={}", userId, e.getMessage(), e);
                return CompletableFuture.completedFuture(false);
            }
        } else {
            log.info("用户不在线，存储离线消息: userId={}", userId);
            storeOfflineMessage(userId, message);
            return CompletableFuture.completedFuture(false);
        }
    }

    public void broadcastMessage(MessageDetailDTO message) {
        if (message == null) {
            return;
        }

        sessions.forEach((userId, session) -> {
            if (session.isOpen()) {
                sendMessage(userId, message);
            } else {
                storeOfflineMessage(userId, message);
            }
        });
        log.info("广播消息完成: messageType={}", message.getType());
    }

    private void sendUnreadMessages(Long userId) {
        try {
            if (messageService == null) {
                log.error("messageService未注入");
                return;
            }

            messageService.getUserMessages(userId, null, new Page<>(1, 100))
                .getRecords()
                .stream()
                .filter(msg -> !Boolean.TRUE.equals(msg.getIsRead()))
                .forEach(msg -> sendMessage(userId, msg));
        } catch (Exception e) {
            log.error("发送未读消息失败: userId={}", userId, e);
        }
    }

    private void sendOfflineMessages(Long userId) {
        ConcurrentLinkedQueue<MessageDetailDTO> messages = OFFLINE_MESSAGES.get(userId);
        if (messages != null && !messages.isEmpty()) {
            MessageDetailDTO message;
            while ((message = messages.poll()) != null) {
                sendMessage(userId, message);
            }
            OFFLINE_MESSAGES.remove(userId);
        }
    }

    private void storeOfflineMessage(Long userId, MessageDetailDTO message) {
        OFFLINE_MESSAGES.computeIfAbsent(userId, k -> new ConcurrentLinkedQueue<>()).offer(message);
        
        // 同时保存到数据库
        try {
            MessageDTO messageDTO = convertToMessageDTO(message);
            messageService.sendMessage(messageDTO);
        } catch (Exception e) {
            log.error("存储离线消息到数据库失败: userId={}", userId, e);
        }
    }

    private void sendHeartbeatResponse(Long userId) {
        try {
            WebSocketSession session = sessions.get(userId);
            if (session != null && session.isOpen()) {
                JsonObject response = new JsonObject();
                response.addProperty("type", "heartbeat");
                response.addProperty("timestamp", System.currentTimeMillis());
                session.sendMessage(new TextMessage(gson.toJson(response)));
            }
        } catch (IOException e) {
            log.error("发送心跳响应失败: userId={}", userId, e);
        }
    }

    private void sendMessageAck(Long userId, String messageId) {
        try {
            WebSocketSession session = sessions.get(userId);
            if (session != null && session.isOpen()) {
                JsonObject response = new JsonObject();
                response.addProperty("type", "ack");
                response.addProperty("messageId", messageId);
                session.sendMessage(new TextMessage(gson.toJson(response)));
            }
        } catch (IOException e) {
            log.error("发送消息确认失败: userId={}, messageId={}", userId, messageId, e);
        }
    }

    private static void checkHeartbeats() {
        LocalDateTime now = LocalDateTime.now();
        LAST_HEARTBEAT.forEach((userId, lastHeartbeat) -> {
            if (ChronoUnit.MILLIS.between(lastHeartbeat, now) > HEARTBEAT_TIMEOUT) {
                log.warn("用户心跳超时，关闭连接: userId={}", userId);
                WebSocketSession session = sessions.get(userId);
                if (session != null) {
                    try {
                        session.close();
                    } catch (IOException e) {
                        log.error("关闭超时连接失败: userId={}", userId, e);
                    }
                }
                sessions.remove(userId);
                LAST_HEARTBEAT.remove(userId);
            }
        });
    }

    private String generateMessageId() {
        return String.format("%d-%d", System.currentTimeMillis(), ThreadLocalRandom.current().nextInt(10000));
    }

    private void scheduleConfirmationTimeout(String messageId, CompletableFuture<Boolean> confirmation) {
        CompletableFuture.runAsync(() -> {
            try {
                Thread.sleep(MESSAGE_CONFIRMATION_TIMEOUT);
                CompletableFuture<Boolean> future = MESSAGE_CONFIRMATIONS.remove(messageId);
                if (future != null && !future.isDone()) {
                    future.complete(false);
                    log.warn("消息确认超时: messageId={}", messageId);
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                log.error("消息确认超时任务被中断: messageId={}", messageId);
            }
        });
    }

    private MessageDTO convertToMessageDTO(MessageDetailDTO detailDTO) {
        MessageDTO dto = new MessageDTO();
        dto.setTitle(detailDTO.getTitle());
        dto.setContent(detailDTO.getContent());
        dto.setType(detailDTO.getType());
        dto.setTargetType(2); // 指定用户
        // 如果 userId 为空，使用 targetIds 中的第一个 ID
        Long userId = detailDTO.getUserId() != null ? detailDTO.getUserId() : 
            (detailDTO.getTargetIds() != null && !detailDTO.getTargetIds().isEmpty() ? 
            detailDTO.getTargetIds().get(0) : null);
        if (userId == null) {
            throw new IllegalArgumentException("无法确定消息接收用户");
        }
        dto.setTargetIds(Collections.singletonList(userId));
        return dto;
    }

    private Long getUserId(WebSocketSession session) {
        try {
            return Long.parseLong(session.getAttributes().get("userId").toString());
        } catch (Exception e) {
            return null;
        }
    }
}